import asyncio
import json
import logging

from langgraph.checkpoint.memory import InMemorySaver
from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent
from my_llm import llm

# 设置记忆存储
checkpointer = InMemorySaver()

# 读取提示词
with open("agent_prompts.txt", "r", encoding="utf-8") as f:
    prompt = f.read()

# 设置对话配置
config = {
    "configurable": {
        "thread_id": "1"
    }
}


# 环境配置
# 读取servers_config.json中的MCP服务器信息
def get_servers_config():
    with open("servers_config.json", "r", encoding="utf-8") as f:
        return json.load(f).get("mcpServers", {})


# 主逻辑
async def run_chat_loop() -> None:
    servers_config = get_servers_config()
    # 连接多台MCP服务器
    mcp_client = MultiServerMCPClient(servers_config)
    # 获取tools，可以用作LangGraph Tools列表
    tools = await mcp_client.get_tools()
    # 构建agent
    agent = create_react_agent(model=llm, tools=tools, prompt=prompt, checkpointer=checkpointer)
    # CLI聊天
    print("\n MCP Agent 已启动，输入quit退出。")
    while True:
        user_input = input("\n你：").strip()
        if user_input.lower() == "quit":
            print("已退出。")
            break
        try:
            res = await agent.ainvoke(
                {"messages": [{"role": "user", "content": user_input}]},
                config
            )
            print(f"\nAI：{res['messages'][-1].content}")
        except Exception as e:
            print(f"\n出错了：{e}")

    # 清理
    # 并没有这个方法
    # await mcp_client.cleanup()
    print("资源已清理，Bye！")


# 入口
if __name__ == "__main__":
    # logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s -%(message)s")
    asyncio.run(run_chat_loop())
